package com.abc.ollama.utils;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

public class OllamaUtils {

    /**
     * 计算两个向量的余弦相似度
     *
     * @param vec1 向量1
     * @param vec2 向量2
     * @return 余弦相似度
     */
    public static double calculate(List<Double> vec1, List<Double> vec2) {
        if (vec1.size() != vec2.size()) {
            System.out.println("vec1.length : " + vec1.size() + " vec2.length : " + vec2.size());
            throw new IllegalArgumentException("The length of two vectors must be equal");
        }
        double dotProduct = 0.0; // 点积
        double normVec1 = 0.0; // 向量1的模
        double normVec2 = 0.0; // 向量2的模

        for (int i = 0; i < vec1.size(); i++) {
            double val1 = vec1.get(i);
            double val2 = vec2.get(i);
            dotProduct += val1 * val2;
            normVec1 += Math.pow(val1, 2);
            normVec2 += Math.pow(val2, 2);
        }

        return dotProduct / (Math.sqrt(normVec1) * Math.sqrt(normVec2));
    }

    /**
     * 将逗号分隔的字符串转换为 List<Double>
     *
     * @param vectorStr 逗号分隔的字符串
     * @return List<Double>
     */
    public static List<Double> parseVector(String vectorStr) {
        if (vectorStr == null || vectorStr.isEmpty()) {
            return List.of();
        }
        return Arrays.stream(vectorStr.split(","))
                .map(String::trim)
                .map(Double::parseDouble)
                .collect(Collectors.toList());
    }
}
